package com.meituan.hive.udf.generic;

import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.exec.UDFArgumentLengthException;
import org.apache.hadoop.hive.ql.exec.UDFArgumentTypeException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.serde2.objectinspector.*;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorUtils;

import java.util.Map;
import java.util.Set;

/**
 * Created by zhaoxiang on 2017/8/9.
 *
 * 根据key的范围对value求和
 * func(map_field, beginKey, endKey, initialVal=0)
 */
public class GenericUDFMapSumValueBykeyRange  extends GenericUDF{
    //private final JavaIntObjectInspector javaIntObjectInspector = PrimitiveObjectInspectorFactory.javaIntObjectInspector;
    private transient MapObjectInspector mapOI;
    private transient PrimitiveObjectInspector mapValueOI;
    private transient PrimitiveObjectInspector mapKeyOI;
    private transient PrimitiveObjectInspector returnOI;
    private transient PrimitiveObjectInspector rangeKeyOI;
    //private transient ObjectInspector[] argumentOIs;
    private transient ObjectInspectorConverters.Converter converter;

    @Override
    public ObjectInspector initialize(ObjectInspector[] arguments)
            throws UDFArgumentException {
        if (arguments.length != 4) {
            throw new UDFArgumentLengthException("The function MapSumValueBykey only accepts 4 argument.");
        } else if (!(arguments[0] instanceof MapObjectInspector)) {
            throw new UDFArgumentTypeException(0, "\""
                    + ObjectInspector.Category.MAP.toString().toLowerCase()
                    + "\" is expected at function MapSumValueBykey, " + "but \""
                    + arguments[0].getTypeName() + "\" is found");
        }
        //argumentOIs = arguments;
        mapOI = (MapObjectInspector) arguments[0];
        mapKeyOI = (PrimitiveObjectInspector)mapOI.getMapKeyObjectInspector();
        mapValueOI = (PrimitiveObjectInspector)mapOI.getMapValueObjectInspector();

        rangeKeyOI = (PrimitiveObjectInspector)arguments[1];
        returnOI = (PrimitiveObjectInspector)arguments[3];

        converter = ObjectInspectorConverters.getConverter(PrimitiveObjectInspectorFactory.javaDoubleObjectInspector, returnOI);

        return ObjectInspectorUtils.getStandardObjectInspector(arguments[3]);
    }

    @Override
    public Object evaluate(GenericUDF.DeferredObject[] arguments) throws HiveException {
        Object mapObj = arguments[0].get();
        double sum = 0;
        if(arguments[3].get() != null)
            sum = PrimitiveObjectInspectorUtils.getDouble(arguments[3].get(), returnOI);

        if(mapObj == null)
            return  converter.convert(sum);

        String beinkeyObj = PrimitiveObjectInspectorUtils.getString(
                arguments[1].get(), rangeKeyOI);
        String endKeyObj = PrimitiveObjectInspectorUtils.getString(
                arguments[2].get(), rangeKeyOI);

        Map<?,?> mapVal = mapOI.getMap(mapObj);
        Set<?> keySet = mapVal.keySet();

        for(Object key: keySet){
            //ObjectInspectorUtils.compare(o1, valInspector, o2, valInspector);
            String strKey = PrimitiveObjectInspectorUtils.getString( key, mapKeyOI);
            if (strKey.compareTo(beinkeyObj) <0  || strKey.compareTo(endKeyObj) > 0)
                continue;
            sum += PrimitiveObjectInspectorUtils.getDouble(
                    mapVal.get(key), mapValueOI);
        }

        return converter.convert(sum);
    }

    @Override
    public String getDisplayString(String[] children) {
        assert children.length == 1;
        return getStandardDisplayString("MapSumValueBykey", children);
    }

}
